import os

import matplotlib.pyplot as plt
import numpy as np
from scipy.optimize import minimize


# ====================================================
def normalize_by_sum(arr, binwidth):
    sum = np.sum(arr)
    if sum == 0:
        return np.zeros_like(arr)
    return arr / sum / binwidth


def weighted_stddev(x, y, eps=1e-12):
    y = np.maximum(y, eps)
    w = y / np.sum(y)
    mu = np.sum(w * x)
    var = np.sum(w * (x - mu) ** 2)
    return np.sqrt(var)


def multi_gauss_model(x, params, n_gauss, with_background=False):
    """
    参数排列: [normFactor, w2..wn, s1..sn, (c)]
    normFactor: 总幅度scale factor
    wi: 各高斯分量权重
    si: 各高斯分量sigma
    c : 常数背景项 (可选)
    """
    normFactor = params[0]
    w = params[1:n_gauss]
    s = params[n_gauss : 2 * n_gauss]
    y_pred = np.zeros_like(x, dtype=float)

    w1 = 1 - np.sum(w)
    y_pred += w1 * np.exp(-0.5 * (x / s[0]) ** 2) / (np.sqrt(2 * np.pi) * s[0])
    for wi, si in zip(w, s[1:]):
        y_pred += wi * np.exp(-0.5 * (x / si) ** 2) / (np.sqrt(2 * np.pi) * si)
    y_pred *= normFactor
    if with_background:
        y_pred += params[-1]
    return y_pred


def sigma_smooth_penalty(s, s_prev, lam_s=1e3):
    return lam_s * np.sum(((s - s_prev) / np.maximum(s_prev, 1e-3)) ** 2)


def mid_tail_weights(x, center=15.0, width=3.0, lam_m=10.0):
    """
    在 |x|位于center附近时放大权重, 来改善这个区域的拟合.
    """
    r = np.abs(x)
    return 1.0 + lam_m * np.exp(-0.5 * ((x - center) / max(width, 1e-6)) ** 2)


def mid_tail_detect(x, y, y_pred0, r_min=5.0, r_max=30.0, smooth_window=10, width_clip=(3.0, 12.0), lam_m_clip=(6.0, 30.0)):
    """
    初步拟合得到 y_pred0，然后根据误差情况自动配置 (center, width, lam_m).
    - 在 [r_min, r_max] 的范围内寻找误差最大的位置;
    - width 为HWHM, 并裁剪到 width_clip；
    - lam_m 与最大相对误差线性映射到 lam_m_clip.
    """
    eps = 1e-12
    r = np.abs(x)
    rel_err = np.abs(y_pred0 - y) / np.maximum(y, eps)

    # 只取主峰两侧的中间部分
    m = (r >= r_min) & (r <= r_max)
    if not np.any(m):
        return dict(center=15.0, width=5.0, lam_m=1.0)

    r_m = r[m]
    rel_err_m = rel_err[m]

    # 因为涉及到正负两侧的range, sort并做简易平滑
    idx = np.argsort(r_m)
    r_s = r_m[idx]
    rel_s = rel_err_m[idx]

    if smooth_window > 1 and smooth_window < rel_s.size:
        k = smooth_window
        kernel = np.ones(k, dtype=float) / k
        rel_sm = np.convolve(rel_s, kernel, mode="same")
    else:
        rel_sm = rel_s

    # 误差最大位置
    i_peak = int(np.argmax(rel_sm))
    center = r_s[i_peak]

    # width
    half = 0.5 * rel_sm[i_peak]
    # 左侧
    iL = 0
    for i in range(i_peak, -1, -1):
        if rel_sm[i] < half:
            iL = i
            break
    # 右侧
    iR = len(rel_sm) - 1
    for i in range(i_peak, len(rel_sm)):
        if rel_sm[i] < half:
            iR = i
            break
    width = 0.5 * max(1e-6, r_s[iR] - r_s[iL])
    width = np.clip(width, width_clip[0], width_clip[1])

    max_rel = rel_sm[i_peak]
    print(f"max_rel = {max_rel:.4f}")
    lam_m = 8.0 + 18.0 * (max_rel / (max_rel + 0.1))
    lam_m = np.clip(lam_m, lam_m_clip[0], lam_m_clip[1])

    return dict(center=center, width=width, lam_m=lam_m)


def mid_tail_detect_v2(x, y, y_pred0, r_min=5.0, r_max=40.0, smooth_window=15, width_clip=(2.0, 40.0), lam_m_clip=(6.0, 40.0), rel_err=0.01, abs_err=0.0):
    """
    查找主峰两侧的中尾部 y_pred < y 的区域.
    """
    eps = 1e-12
    diff = y_pred0 - y
    thres = np.maximum(abs_err, rel_err * np.maximum(y, eps))

    # y_pred < y 部分
    sel_mask = (diff < -thres) & (np.abs(x) >= r_min) & (np.abs(x) <= r_max)
    if not np.any(sel_mask):
        return dict(x_left=None, x_right=None, center=15.0, width=5.0, lam_m=1.0)

    x_sel = x[sel_mask]
    rel_sel = -(diff[sel_mask]) / np.maximum(y[sel_mask], eps)

    idx = np.argsort(x_sel)
    x_s = x_sel[idx]
    rel_s = rel_sel[idx]

    # 找出最长的连续 y_pred < y 部分
    if len(x_s) < 3:
        return dict(x_left=x_s.min(), x_right=x_s.max(), center=np.mean(x_s), width=2.0, lam_m=8.0)

    dx = 0.1  # 相邻点横坐标间距
    gap_thres = 1.5 * dx

    # 如果相邻点间距大于gap_thres则认为不连续
    starts = [0]
    for i in range(1, len(x_s)):
        if (x_s[i] - x_s[i - 1]) > gap_thres:
            starts.append(i)
    ends = starts[1:] + [len(x_s)]

    seg_lengths = [ends[k] - starts[k] for k in range(len(starts))]
    k_max = int(np.argmax(seg_lengths))
    left, right = starts[k_max], ends[k_max]
    x_seg, rel_seg = x_s[left:right], rel_s[left:right]

    if len(x_seg) < 3:
        return dict(x_left=x_seg.min(), x_right=x_seg.max(), center=np.mean(x_seg), width=2.0, lam_m=8.0)

    # 简单平滑并确定误差最大位置
    if smooth_window > 1 and smooth_window < len(rel_seg):
        kernel = np.ones(smooth_window) / smooth_window
        rel_sm = np.convolve(rel_seg, kernel, mode="same")
    else:
        rel_sm = rel_seg

    # 找y_pred < y 部分的center的位置
    i_peak = np.argmax(rel_sm)
    x_peak = x_seg[i_peak]
    peak_val = rel_sm[i_peak]
    print(f"peak_val: {peak_val:.3f}")

    # 定位width
    half = 0.1 * peak_val
    iL, iR = 0, len(x_seg) - 1
    for i in range(i_peak, -1, -1):
        if rel_sm[i] < half:
            iL = i
            break
    for i in range(i_peak, len(x_seg)):
        if rel_sm[i] < half:
            iR = i
            break

    x_left, x_right = x_seg[iL], x_seg[iR]
    width = 0.5 * abs(x_right - x_left)
    width = np.clip(width, width_clip[0], width_clip[1])

    # lam_m = 8.0 + 18.0 * (peak_val / (peak_val + 0.2))
    lam_m = 8.0 + 30.0 * (1 - np.exp(-3.0 * peak_val))
    lam_m = np.clip(lam_m, lam_m_clip[0], lam_m_clip[1])

    return dict(center=x_peak, width=width, lam_m=lam_m)


def objective_log(params, x, y, n_gauss, with_background, lam_n=1e3, s_prev=None, mid_tail_weight_para=None):
    """
    第一阶段目标：log residual平方和 + 权重和的软约束 λ * (sum(w)-1)^2
    """
    if np.any(params < 0):
        return 1e20
    s = params[n_gauss : 2 * n_gauss]
    if np.any(s <= 0):
        return 1e20
    y_pred = multi_gauss_model(x, params, n_gauss, with_background)
    eps = 1e-12
    if mid_tail_weight_para is None:
        W = 1.0
    else:
        W = mid_tail_weights(x, **mid_tail_weight_para)
    # weights = np.maximum(y, eps)
    # loss = np.sum(weights * (np.log(y_pred+eps) - np.log(y+eps))**2)
    loss = np.sum(W * (np.log(y_pred + eps) - np.log(y + eps)) ** 2)
    # 软约束：权重和接近1
    #    w_sum = np.sum(params[1:n_gauss+1])
    #    loss += lam_w * (w_sum - 1.0)**2
    # 软约束：normFactor接近1
    normFactor = params[0]
    loss += lam_n * (normFactor - 1.0) ** 2
    if s_prev is not None:
        loss += sigma_smooth_penalty(s, s_prev)
    return loss


def objective_linear(params, x, y, n_gauss, with_background, lam_n=1e3, s_prev=None, mid_tail_weight_para=None):
    """
    第二阶段目标：linear residual平方和 + 权重和软约束 λ * (sum(w)-1)^2
    """
    if np.any(params < 0):
        return 1e20
    s = params[n_gauss : 2 * n_gauss]
    if np.any(s <= 0):
        return 1e20
    y_pred = multi_gauss_model(x, params, n_gauss, with_background)
    if mid_tail_weight_para is None:
        W = 1.0
    else:
        W = mid_tail_weights(x, **mid_tail_weight_para)
    loss = W * np.sum((y_pred - y) ** 2)
    # 软约束：权重和接近1
    #    w_sum = np.sum(params[1:n_gauss+1])
    #    loss += lam_w * (w_sum - 1.0)**2
    # 软约束：normFactor接近1
    normFactor = params[0]
    loss += lam_n * (normFactor - 1.0) ** 2
    if s_prev is not None:
        loss += sigma_smooth_penalty(s, s_prev)
    return loss


def build_constraints(n_gauss, is_post_bragg=False):
    cons = []
    # 不等式约束：sum(weight)<1, sigma递增
    # for i in range(n_gauss-1):
    #    cons.append({"type": "ineq", "fun": lambda p, i=i: p[i+1] - p[i+2]})
    cons.append({"type": "ineq", "fun": lambda p, n=n_gauss: 1 - np.sum(p[1:n])})

    for i in range(n_gauss - 1):
        if is_post_bragg and i == 0:
            continue
        cons.append({"type": "ineq", "fun": lambda p, i=i, n=n_gauss: p[n + i + 1] - p[n + i]})
    return cons


def build_bounds(x, n_gauss, s_prev=None, step_range=0.2, is_post_bragg=False, step_range_relax=5):
    lower = [0.5] + [1e-8] * (n_gauss - 1) + [1e-3] * n_gauss
    upper = [1.5] + [1.0] * (n_gauss - 1) + [(x.max() - x.min())] * n_gauss

    # 对于相邻层的拟合, 限制参数变化范围
    if s_prev is not None:
        for j in range(n_gauss):
            idx = 1 + (n_gauss - 1) + j
            # 判断是否是第一高斯sigma, 且在Bragg峰后
            # if j == 0 and is_post_bragg:
            if is_post_bragg:
                s_range = step_range_relax
            else:
                s_range = step_range

            lb = max(1e-3, (1 - s_range) * s_prev[j])
            ub = max(lb * 1.05, (1 + s_range) * s_prev[j])
            lower[idx] = lb
            upper[idx] = ub
    print(list(zip(lower, upper)))
    return list(zip(lower, upper))


def fit_multi_gauss(x, y, n_gauss=2, init_params=None, with_background=False, is_post_bragg=False, s_prev=None, lam_n=1e3, info_verbose=False, info_tag=None):
    x = np.asarray(x, float)
    y = np.asarray(y, float)

    # 初值
    if init_params is None:
        norm_init = 1.0
        w_init = np.ones(n_gauss - 1) * 0.5 / (n_gauss - 1)  # 平均分配权重
        # TODO weighted_stddev在不是全范围时可能有问题?
        s_guess = weighted_stddev(x, y)
        s_init = [s_guess * (0.8 + 2 * i) for i in range(n_gauss)]
        # s_init = np.logspace(np.log10(1e-2), np.log10(30), n_gauss)
        p0 = np.array([norm_init] + list(w_init) + list(s_init) + ([y.min()] if with_background else []))
        print("初始param: ", p0)
    else:
        p0 = init_params.copy()

    #    # 边界
    #    lower = [0.5] + [1e-8]*(n_gauss-1) + [1e-3]*n_gauss
    #    upper = [1.5] + [1.0]*(n_gauss-1) + [5*(x.max()-x.min())]*n_gauss
    #    if with_background:
    #        lower.append(0)
    #        upper.append(0.5)
    #    bounds = list(zip(lower, upper))
    bounds = build_bounds(x, n_gauss, s_prev, is_post_bragg=is_post_bragg)

    constraints = build_constraints(n_gauss, is_post_bragg)

    # 预拟合, 用来检测主峰两侧的误差, 以设置mid_tail_weight
    res0 = minimize(objective_log, p0, args=(x, y, n_gauss, with_background, lam_n, s_prev, None), method="SLSQP", bounds=bounds, constraints=constraints, options={"maxiter": 500, "ftol": 1e-10})
    y_pred0 = multi_gauss_model(x, res0.x, n_gauss, with_background)

    mid_tail_weight_para = mid_tail_detect_v2(x, y, y_pred0)
    # print("mid_tail_weight_para: ", mid_tail_weight_para)
    if info_verbose:
        prefix = f"[{info_tag}] " if info_tag is not None else ""
        print(f"{prefix}mid_tail_weight_para:  center={mid_tail_weight_para['center']:.3f}, width={mid_tail_weight_para['width']:.3f}, lam_m={mid_tail_weight_para['lam_m']:.1f}, n={n_gauss}")

    # if n_gauss == 3:
    #    mid_tail_weight_para["lam_m"] = float(np.clip(1.5 * mid_tail_weight_para["lam_m"], 6.0, 40.0))

    # 第一阶段：log residual拟合, 更关注尾部, 避免被主峰主导
    res1 = minimize(objective_log, p0, args=(x, y, n_gauss, with_background, lam_n), method="SLSQP", bounds=bounds, constraints=constraints, options={"maxiter": 3000, "ftol": 1e-10})

    #    idx_tail_range = np.where(np.abs(x) > 3*res1.x[n_gauss+1])[0]

    #    res1 = minimize(objective_log, res1.x,
    #                      args=(x[idx_tail_range], y[idx_tail_range], n_gauss, with_background, lam_w, lam_n),
    #                      method="SLSQP", bounds=bounds, constraints=constraints,
    #                      options={"maxiter": 2000, "ftol": 1e-10})

    # 第二阶段：linear residual精修, 调整整体幅度
    res2 = minimize(objective_linear, res1.x, args=(x, y, n_gauss, with_background, lam_n), method="SLSQP", bounds=bounds, constraints=constraints, options={"maxiter": 2000, "ftol": 1e-10})

    if not res2.success:
        print("优化未收敛：", res2.message)

    return res2.x, res2


################################
# 开始拟合
#
################################

# 创建输出目录
enedir = "275MeV"
output_dir = f"fit_plots_{enedir}_5g_midtail_weight_v2"
os.makedirs(f"{output_dir}", exist_ok=True)
fit_results_txt = f"fit_results_{enedir}_5g_midtail_weight_v2.txt"


# prepare input data
Zbin = 1000
Zmax = 50.0
Ybin = 1000
Ymax = 50
Ymin = -50
Ymax_fit = 50
Ymin_fit = -50

Zaxis = np.arange(Zmax / Zbin / 2.0, Zmax, Zmax / Zbin)
# print('Zaxis: {}'.format(Zaxis))
Yaxis_raw = np.arange(Ymin + (Ymax - Ymin) / Ybin / 2, Ymax, (Ymax - Ymin) / Ybin)

#     _2            _1              _3
# BLOOD_ICRP, BONE_COMPACT_ICRU, BRAIN_ICRP
# material = 'BLOOD_ICRP'
# material = 'BONE_COMPACT_ICRU'
# material = 'BRAIN_ICRP'
material = "WATER"
parname = "he"
# enedir = '125MeV'
dirtag = "filter"


# read .bin file and restore array shape
path = "/data/wenxiao/cmcvalidate_letd_new/cudaCMC_validate_Geant4/G4Project/" + parname + "_" + enedir + "/"

arr_he = np.fromfile(path + "eDep.bin", dtype=np.float64).reshape(Zbin, Ybin)

IDD = arr_he.sum(axis=1)


# N-高斯拟合设置
n_gauss = 5
results = np.zeros((Zbin, 2 * n_gauss + 1))  # 存储所有深度的拟合结果


# 剂量最大位置
idz0 = np.argmax(arr_he.sum(axis=1))
peak_idd = IDD[idz0]

# find peak first?
# z_depth_index = int(idz0)
z_depth_index = int(0.85 * idz0)  # 选一个peak之前的位置
# z_depth_index = 78

print("Zaxis[{}]: {}".format(z_depth_index, Zaxis[z_depth_index]))

# 归一化数据
profile = arr_he[z_depth_index, :]
profile = normalize_by_sum(profile, (Ymax - Ymin) / Ybin)


# ===================================
# 限制拟合数据的横坐标范围
mask = (Yaxis_raw >= Ymin_fit) & (Yaxis_raw <= Ymax_fit)
Yaxis = Yaxis_raw[mask]
profile = profile[mask]

# ===================================


# 先拟合布拉格峰位附近
x, y = Yaxis, profile
params, res = fit_multi_gauss(x, y, n_gauss, info_verbose=True, info_tag=f"idx={z_depth_index}")
print("params", params)
results[z_depth_index] = np.hstack(([params[0]], [1 - np.sum(params[1:n_gauss])], params[1:n_gauss], params[n_gauss : 2 * n_gauss]))
print("normFactor:", params[0], "第一高斯权重", 1 - np.sum(params[1:n_gauss]), "其余拟合权重:", params[1:n_gauss], "sigma:", params[n_gauss : 2 * n_gauss], "其余权重和:", np.sum(params[1:n_gauss]))
y_fit = multi_gauss_model(x, params, n_gauss, False)


y_min = np.min(profile)
y_max = np.max(profile)
y_lower = max(y_min * 0.5, 1e-12)
y_upper = y_max * 2
plt.figure(dpi=150)
plt.semilogy(x, profile, "o", label="data", markersize=4, markerfacecolor="none", alpha=0.8)
plt.semilogy(x, y_fit, "-", label="fit")
plt.legend()
plt.xlabel("x [mm]")
plt.ylabel("Dose [a.u.]")
plt.title(f"Depth idx={z_depth_index} (z={Zaxis[z_depth_index]:.3f} cm)")
plt.ylim(y_lower, y_upper)
plt.savefig(f"{output_dir}/depth_{z_depth_index:03d}.png")
# plt.savefig(f"depth_{z_depth_index:03d}.png")
plt.close()


# 向更浅深度递推拟合
for i in range(z_depth_index - 1, -1, -1):
    profile = normalize_by_sum(arr_he[i, :], (Ymax - Ymin) / Ybin)
    profile = profile[mask]
    s_prev = params[n_gauss : 2 * n_gauss]
    params, _ = fit_multi_gauss(x, profile, n_gauss, init_params=params, s_prev=s_prev, info_verbose=True, info_tag=f"idx={i}")
    print("params", params)
    results[i] = np.hstack(([params[0]], [1 - np.sum(params[1:n_gauss])], params[1:n_gauss], params[n_gauss : 2 * n_gauss]))
    print("normFactor:", params[0], "第一高斯权重", 1 - np.sum(params[1:n_gauss]), "其余拟合权重:", params[1:n_gauss], "sigma:", params[n_gauss : 2 * n_gauss], "其余权重和:", np.sum(params[1:n_gauss]))
    y_fit = multi_gauss_model(x, params, n_gauss, False)

    y_min = np.min(profile)
    y_max = np.max(profile)
    y_lower = max(y_min * 0.5, 1e-12)
    y_upper = y_max * 2
    # 绘图并保存
    plt.figure(dpi=150)
    plt.semilogy(x, profile, "o", label="data", markersize=4, markerfacecolor="none", alpha=0.8)
    plt.semilogy(x, y_fit, "-", label="fit")
    plt.legend()
    plt.xlabel("x [mm]")
    plt.ylabel("Dose [a.u.]")
    plt.title(f"Depth idx={i} (z={Zaxis[i]:.3f} cm)")
    plt.ylim(y_lower, y_upper)
    plt.savefig(f"{output_dir}/depth_{i:03d}.png")
    plt.close()

# 重置params
# params = results[z_depth_index]
params = np.delete(results[z_depth_index], 1)

# 向更深深度递推拟合, 有深度的截止条件
# for i in range(z_depth_index+1, min(Zbin, 2*z_depth_index, z_depth_index+40)):
for i in range(z_depth_index + 1, min(Zbin, idz0 + 40)):
    if IDD[i] < 0.01 * peak_idd:
        print(f"深度 {i} 已低于峰值 1%，停止拟合。")
        break

    profile = normalize_by_sum(arr_he[i, :], (Ymax - Ymin) / Ybin)
    profile = profile[mask]
    s_prev = params[n_gauss : 2 * n_gauss]
    if i > int(idz0):
        params, _ = fit_multi_gauss(Yaxis, profile, n_gauss, init_params=params, is_post_bragg=True, s_prev=s_prev, info_verbose=True, info_tag=f"idx={i}")
    else:
        params, _ = fit_multi_gauss(Yaxis, profile, n_gauss, init_params=params, is_post_bragg=False, s_prev=s_prev, info_verbose=True, info_tag=f"idx={i}")
    print("params", params)
    results[i] = np.hstack(([params[0]], [1 - np.sum(params[1:n_gauss])], params[1:n_gauss], params[n_gauss : 2 * n_gauss]))
    print("normFactor:", params[0], "第一高斯权重", 1 - np.sum(params[1:n_gauss]), "其余拟合权重:", params[1:n_gauss], "sigma:", params[n_gauss : 2 * n_gauss], "其余权重和:", np.sum(params[1:n_gauss]))
    y_fit = multi_gauss_model(x, params, n_gauss, False)

    y_min = np.min(profile)
    y_max = np.max(profile)
    y_lower = max(y_min * 0.5, 1e-12)
    y_upper = y_max * 2
    # 绘图并保存
    plt.figure(dpi=150)
    plt.semilogy(x, profile, "o", label="data", markersize=4, markerfacecolor="none", alpha=0.8)
    plt.semilogy(x, y_fit, "-", label="fit")
    plt.legend()
    plt.xlabel("x [mm]")
    plt.ylabel("Dose [a.u.]")
    plt.title(f"Depth idx={i} (z={Zaxis[i]:.3f} cm)")
    plt.ylim(y_lower, y_upper)
    plt.savefig(f"{output_dir}/depth_{i:03d}.png")
    plt.close()


# 保存结果到txt
np.savetxt(f"{fit_results_txt}", results, header="norm, w1..wn, sigma1..n", fmt="%.6e")
print(f"拟合完成，结果已保存到 {fit_results_txt} 和 {output_dir}/*.png")
